#include "cem/assembler.h"
#include "cem/boundary_condition.h"
#include "cem/constants.h"
#include "cem/electromagnetics_system.h"
#include "cem/material_database.h"
#include "cem/mesh_manager.h"
#include "cem/parameter_reader.h"
#include "cem/post_processor.h"
#include "cem/solver_manager.h"
#include "cem/timer.h"

#include <libmesh/boundary_info.h>
#include <libmesh/elem.h>
#include <libmesh/libmesh.h>
#include <libmesh/mesh.h>
#include <libmesh/parallel.h>
#include <mpi.h>

#include <algorithm>
#include <cmath>
#include <complex>
#include <iostream>
#include <optional>
#include <string_view>
#include <vector>

// 函数: print_banner
// 作用: 打印程序与并行环境的基本信息，并按 rank 顺序串行输出初始化行
// 参数:
//   init: libMesh 初始化对象引用，用于获取通信器等
// 返回: 无
static void
print_banner(const libMesh::LibMeshInit & init)
{
    const auto & comm = init.comm();
    const int world_size = comm.size();

    std::cout << world_size << " MPI processor initialized." << std::endl;
}

// 结构体: FilteredArgs
// 作用: 承载为 libMesh 过滤后的命令行参数副本（仅用于传入 LibMeshInit）
// 说明: 只删除 --mesh/--config 及其取值（或等号形式），其余参数保持不变
struct FilteredArgs
{
    int argc;                 // 过滤后参数个数
    std::vector<char *> argv; // 指向原始 argv 中对应元素的指针数组
};

// 函数: filter_argv_for_libmesh
// 作用: 基于原始 argc/argv 生成供 libMesh 使用的“已过滤”参数数组
// 规则:
//   - 删除: "--mesh <path>", "--config <path>", "--mesh=<path>", "--config=<path>"
//   - 保留: 其他所有参数（含程序名 argv[0]）
// 参数:
//   argc: 原始参数数量
//   argv: 原始参数指针数组
// 返回: FilteredArgs, 可用于构造 LibMeshInit
static FilteredArgs
filter_argv_for_libmesh(int argc, char ** argv)
{
    FilteredArgs out;
    out.argv.reserve(static_cast<size_t>(argc) + 1); // 预留并保证可追加 nullptr 终止符

    auto starts_with = [](std::string_view s, std::string_view prefix)
    { return s.size() >= prefix.size() && s.compare(0, prefix.size(), prefix) == 0; };

    for (int i = 0; i < argc; ++i)
    {
        std::string_view tok(argv[i] ? argv[i] : "");
        if (i == 0)
        {
            // 保留程序名
            out.argv.push_back(argv[i]);
            continue;
        }

        // 处理等号形式: --mesh=..., --config=...
        if (starts_with(tok, "--mesh=") || starts_with(tok, "--config="))
        {
            continue; // 丢弃整个 token
        }

        // 处理分离形式: --mesh <value> 或 --config <value>
        if (tok == "--mesh" || tok == "--config")
        {
            // 跳过当前选项以及其后的一个取值（若存在）
            if (i + 1 < argc)
                ++i;
            continue;
        }

        // 其余 token 原样保留
        out.argv.push_back(argv[i]);
    }

    // 重要: 为传给 MPI/LibMesh 的 argv 添加 NULL 终止符，避免第三方库按 NULL 结尾遍历时越界
    out.argv.push_back(nullptr);
    // argc 统计不包含末尾的 nullptr
    out.argc = static_cast<int>(out.argv.size()) - 1;
    return out;
}

// 函数: build_second_polarization_specs
// 作用: 基于第一套边界规格自动生成第二套极化（默认绕 Z 轴 +90° 旋转：x->y, y->-x）。
// 说明:
//   - 保持边界类型、id、不改变幅值/复数模式。
//   - 若原规格未显式提供方向(has_direction=false)，则为其启用方向并设为全局 y 轴。
// 参数:
//   specs_pol1: 第1套极化的边界规格
// 返回:
//   生成的第2套极化边界规格
static std::vector<cem::BoundarySpec>
build_second_polarization_specs(const std::vector<cem::BoundarySpec> & specs_pol1)
{
    std::vector<cem::BoundarySpec> specs_pol2;
    specs_pol2.reserve(specs_pol1.size());

    for (const auto & s1 : specs_pol1)
    {
        cem::BoundarySpec s2 = s1; // 复制类型/幅值/模式等
        if (s1.has_direction)
        {
            // 绕 Z 轴旋转 +90°: (x, y, z) -> (-y, x, z)
            const double x = s1.direction[0];
            const double y = s1.direction[1];
            const double z = s1.direction[2];
            s2.direction = {-y, x, z};
            s2.has_direction = true;
        }
        else
        {
            // 若未提供方向，则启用并设为全局 y 轴
            s2.direction = {0.0, 1.0, 0.0};
            s2.has_direction = true;
        }
        specs_pol2.push_back(s2);
    }

    return specs_pol2;
}

// 函数: allgather_union_int
// 作用: 在指定通信器上对每个进程的本地整数 id 集合执行 Allgatherv，
//       汇总为全局并集（去重升序），并将结果返回到该通信器的所有进程。
// 参数:
//   local_ids - 本地整数 id 集合
//   comm      - libMesh 通信器，通常为子通信器
// 返回: 全局并集的整数 id 列表（去重升序）
static std::vector<int>
allgather_union_int(const std::vector<int> & local_ids,
                    const libMesh::Parallel::Communicator & comm)
{
    int comm_size = 1;
    int comm_rank = 0;
    MPI_Comm mpi_comm = comm.get();
    MPI_Comm_size(mpi_comm, &comm_size);
    MPI_Comm_rank(mpi_comm, &comm_rank);

    const int local_count = static_cast<int>(local_ids.size());
    std::vector<int> counts(comm_size, 0);
    MPI_Allgather(&local_count, 1, MPI_INT, counts.data(), 1, MPI_INT, mpi_comm);

    std::vector<int> displs(comm_size, 0);
    int total = 0;
    for (int i = 0; i < comm_size; ++i)
    {
        displs[i] = total;
        total += counts[i];
    }

    std::vector<int> all_ids(total);
    if (total > 0)
    {
        MPI_Allgatherv(local_ids.data(),
                       local_count,
                       MPI_INT,
                       all_ids.data(),
                       counts.data(),
                       displs.data(),
                       MPI_INT,
                       mpi_comm);
        std::sort(all_ids.begin(), all_ids.end());
        all_ids.erase(std::unique(all_ids.begin(), all_ids.end()), all_ids.end());
    }
    return all_ids;
}

// 新增函数: broadcast_and_sync_boundary_specs
// 作用: 以子通信器 rank 0 的边界规格集合为准，通过 MPI_Bcast 将规格广播到该通信器所有进程，
//       并用广播结果覆盖本地 specs，确保通信器内一致。
// 参数:
//   specs - 本地边界规格集合（入参在各进程上可不一致），广播后统一为子通信器 rank0 的集合
//   comm  - libMesh 通信器，通常为子通信器
static void
broadcast_and_sync_boundary_specs(std::vector<cem::BoundarySpec> & specs,
                                  const libMesh::Parallel::Communicator & comm)
{
    const int root = 0;
    MPI_Comm mpi_comm = comm.get();

    int n = static_cast<int>(specs.size());
    MPI_Bcast(&n, 1, MPI_INT, root, mpi_comm);

    // 每条规格的打包大小
    const int ints_per = 4;    // boundary_id, type, has_direction, mode
    const int doubles_per = 6; // value, dir[3], value_real, value_imag

    std::vector<int> ibuf;
    std::vector<double> dbuf;
    if (comm.rank() == root)
    {
        ibuf.resize(n * ints_per);
        dbuf.resize(n * doubles_per);
        for (int i = 0; i < n; ++i)
        {
            const auto & s = specs[i];
            ibuf[i * ints_per + 0] = s.boundary_id;
            ibuf[i * ints_per + 1] = static_cast<int>(s.type);
            ibuf[i * ints_per + 2] = s.has_direction ? 1 : 0;
            ibuf[i * ints_per + 3] = static_cast<int>(s.mode);

            dbuf[i * doubles_per + 0] = s.value;
            dbuf[i * doubles_per + 1] = s.direction[0];
            dbuf[i * doubles_per + 2] = s.direction[1];
            dbuf[i * doubles_per + 3] = s.direction[2];
            dbuf[i * doubles_per + 4] = s.value_real;
            dbuf[i * doubles_per + 5] = s.value_imag;
        }
    }
    else
    {
        ibuf.resize(n * ints_per);
        dbuf.resize(n * doubles_per);
    }

    if (!ibuf.empty())
        MPI_Bcast(ibuf.data(), static_cast<int>(ibuf.size()), MPI_INT, root, mpi_comm);
    if (!dbuf.empty())
        MPI_Bcast(dbuf.data(), static_cast<int>(dbuf.size()), MPI_DOUBLE, root, mpi_comm);

    specs.resize(n);
    for (int i = 0; i < n; ++i)
    {
        cem::BoundarySpec s;
        s.boundary_id = ibuf[i * ints_per + 0];
        s.type = static_cast<cem::BoundaryType>(ibuf[i * ints_per + 1]);
        s.has_direction = (ibuf[i * ints_per + 2] != 0);
        s.mode = static_cast<cem::BoundaryComplexMode>(ibuf[i * ints_per + 3]);

        s.value = dbuf[i * doubles_per + 0];
        s.direction[0] = dbuf[i * doubles_per + 1];
        s.direction[1] = dbuf[i * doubles_per + 2];
        s.direction[2] = dbuf[i * doubles_per + 3];
        s.value_real = dbuf[i * doubles_per + 4];
        s.value_imag = dbuf[i * doubles_per + 5];

        specs[i] = s;
    }
}
int
main(int argc, char ** argv)
{
    // 对整个 main 过程进行粗粒度计时
    cem::TimerScope scope_main("main");

    // 在初始化 libMesh 之前过滤掉仅供本程序消费的选项，避免 libMesh/GetPot 报告未使用参数
    FilteredArgs filtered = filter_argv_for_libmesh(argc, argv);
    int argc_for_libmesh = filtered.argc;
    char ** argv_for_libmesh = filtered.argv.data();

    libMesh::LibMeshInit init(argc_for_libmesh, argv_for_libmesh);

    print_banner(init);

    // 1) 解析 CLI（使用原始 argc/argv）
    cem::ParameterReader params;
    params.parse_cli(argc, argv);

    // 2) 若给出 --config 则加载 JSON
    if (auto cfg = params.get_string("config"); cfg.has_value())
    {
        if (!params.load_from_file(cfg.value()))
            std::cerr << "Failed to load config: " << cfg.value() << std::endl;
    }

    // === 子通信器分组（基于 N_GROUP，按块分配频率） ===
    const auto & freq_list = params.get_frequencies();
    int n_freq = static_cast<int>(freq_list.size());
    int world_rank = init.comm().rank();
    int world_size = init.comm().size();

    // 读取 N_GROUP; 若未提供则退化为 min(n_freq, world_size) 或 1
    int desired_groups =
        params.get_num_groups().value_or((n_freq > 0) ? std::min(n_freq, world_size) : 1);
    // 为稳健性进行夹取: 不超过 world_size, 也不超过 n_freq(若为 0 则保持为 1)
    int n_groups = std::max(
        1, std::min(desired_groups, (n_freq > 0) ? std::min(n_freq, world_size) : world_size));
    int color = (n_freq > 0) ? (world_rank % n_groups) : 0;

    libMesh::Parallel::Communicator group_comm;
    init.comm().split(color, /*key=*/world_rank, group_comm);
    int group_rank = group_comm.rank();

    // 计算当前 color 组负责的频率索引集合（分块分配）
    std::vector<int> freq_indices;
    if (n_freq > 0)
    {
        int base = n_freq / n_groups; // 每组的基础块大小
        int rem = n_freq % n_groups;  // 余数，前 rem 个组多分 1 个
        int start = color * base + std::min(color, rem);
        int count = base + (color < rem ? 1 : 0);
        for (int k = 0; k < count; ++k)
        {
            freq_indices.push_back(start + k);
        }
    }

    // 2.1) 组分配回显: 在组内 rank0 输出该组的频率索引与实际频率值，便于核查并行分配是否符合预期
    if (group_rank == 0)
    {
        std::cout << "[color=" << color << "][world_rank=" << world_rank
                  << "][group_rank=" << group_rank << "] "
                  << "[group] n_groups=" << n_groups << ", group_size=" << group_comm.size()
                  << ", assigned_freq_indices:";
        for (int idx : freq_indices)
            std::cout << ' ' << idx;
        std::cout << "; assigned_freq_values:";
        for (int idx : freq_indices)
            std::cout << ' ' << freq_list[idx];
        std::cout << std::endl;
    }

    // 3) 读取网格（若未提供则用 example/case1/3D.msh 作为默认）
    std::string mesh_path = "example/case1/3D.msh";
    if (auto m = params.get_string("mesh"); m.has_value())
        mesh_path = m.value();

    // 使用子通信器构造 MeshManager
    cem::MeshManager mesh_manager(group_comm);
    {
        cem::TimerScope scope_mesh("load_mesh");
        mesh_manager.load_gmsh_mesh(mesh_path);
    }

    // 4) 提取并回显区域与边界 id（并行下做通信器内全局并集）
    auto area_ids_local = mesh_manager.get_subdomain_ids();
    auto bnd_ids_local = mesh_manager.get_boundary_ids();

    // 将 libMesh 的 id 类型转换为 int，便于 MPI 收集
    std::vector<int> area_ids_local_i;
    area_ids_local_i.reserve(area_ids_local.size());
    for (auto id : area_ids_local)
        area_ids_local_i.push_back(static_cast<int>(id));

    std::vector<int> bnd_ids_local_i;
    bnd_ids_local_i.reserve(bnd_ids_local.size());
    for (auto id : bnd_ids_local)
        bnd_ids_local_i.push_back(static_cast<int>(id));

    // 在子通信器上做全局并集（所有组内进程得到一致集合）
    std::vector<int> area_ids_union_i = allgather_union_int(area_ids_local_i, group_comm);
    std::vector<int> bnd_ids_union_i = allgather_union_int(bnd_ids_local_i, group_comm);

    // 转回 libMesh 的 id 类型
    std::vector<libMesh::subdomain_id_type> area_ids;
    area_ids.reserve(area_ids_union_i.size());
    for (int id : area_ids_union_i)
        area_ids.push_back(static_cast<libMesh::subdomain_id_type>(id));

    std::vector<libMesh::boundary_id_type> bnd_ids;
    bnd_ids.reserve(bnd_ids_union_i.size());
    for (int id : bnd_ids_union_i)
        bnd_ids.push_back(static_cast<libMesh::boundary_id_type>(id));

    if (group_rank == 0)
    {
        std::cout << "[color=" << color << "][world_rank=" << world_rank
                  << "][group_rank=" << group_rank << "] "
                  << "[mesh] global area ids:";
        for (auto id : area_ids)
            std::cout << ' ' << id;
        std::cout << std::endl;
        std::cout << "[color=" << color << "][world_rank=" << world_rank
                  << "][group_rank=" << group_rank << "] "
                  << "[mesh] global boundary ids:";
        for (auto id : bnd_ids)
            std::cout << ' ' << id;
        std::cout << std::endl;
    }

    // 5) 基于 JSON 的边界组/边界条件
    cem::BoundaryConditionManager bc_manager;
    const auto & json_specs = params.boundary_specs();
    if (!json_specs.empty())
    {
        // 并行场景下，各 rank 所见的本地边界 id 可能不同，若在此按 bnd_ids 过滤，
        // 会导致 rank 0 丢弃其他 rank 上真实存在的边界（例如 Outer）。
        // 为确保全局一致性，这里直接接纳 JSON 中的全部边界条件，
        // 在装配阶段仅对匹配到的边界 id 生效，未匹配的自然忽略。
        for (const auto & s : json_specs)
            bc_manager.add_spec(s);

        if (group_rank == 0)
        {
            std::cout << "[color=" << color << "][world_rank=" << world_rank
                      << "][group_rank=" << group_rank << "] "
                      << "BC Loaded " << bc_manager.specs().size()
                      << " JSON boundary conditions; applying on matching boundary ids."
                      << std::endl;
        }
    }

    // 若提供了 JSON 但由于边界 id 不一致导致未匹配上，且网格仅有单一边界 id，
    // 则将 JSON 中的所有边界条件自动映射到该唯一的网格边界 id，以提升兼容性。
    if (false && bc_manager.specs().empty() && !json_specs.empty())
    {
        if (bnd_ids.size() == 1)
        {
            const int only_id = static_cast<int>(bnd_ids.front());
            for (const auto & s : json_specs)
            {
                cem::BoundarySpec mapped = s;
                mapped.boundary_id = only_id;
                bc_manager.add_spec(mapped);
            }
            if (group_rank == 0)
            {
                std::cout << "BC No matching ids found; auto-mapped all JSON conditions to "
                             "the only mesh boundary id="
                          << only_id << std::endl;
            }
        }
    }

    if (bc_manager.specs().empty())
    {
        // 默认：所有边界齐次 Dirichlet，保证系统可解（使用通信器并集 bnd_ids）
        for (auto id : bnd_ids)
        {
            cem::BoundarySpec spec;
            spec.boundary_id = static_cast<int>(id);
            spec.type = cem::BoundaryType::Dirichlet;
            spec.value = 0.0;
            bc_manager.add_spec(spec);
            if (group_rank == 0)
                std::cout << "[bc] Dirichlet bid=" << id << ", value=0" << std::endl;
        }
    }

    // 同步边界条件集合：以组内 rank0 的集合为准广播到所有进程
    {
        std::vector<cem::BoundarySpec> specs_synced = bc_manager.specs();
        broadcast_and_sync_boundary_specs(specs_synced, group_comm);
        cem::BoundaryConditionManager bc_synced;
        for (const auto & s : specs_synced)
            bc_synced.add_spec(s);
        bc_manager = std::move(bc_synced);
    }

    // 6) 构建材料数据库（占位，当前装配仅使用区域 MU/EPSL/COND 映射）
    cem::MaterialDatabase material_db;

    // 7) 初始化电磁系统
    cem::ElectromagneticsSystem em_system;
    {
        cem::TimerScope scope_init("initialize_system");
        em_system.initialize(params, mesh_manager, material_db, bc_manager);
    }

    // 8) 求解器配置
    cem::SolverManager solver_manager;
    solver_manager.configure_from_parameters(params);

    // 准备两套极化规格：第1套来自初始化时的集合；第2套由规则生成
    std::vector<cem::BoundarySpec> specs_pol1 = bc_manager.specs();
    std::vector<cem::BoundarySpec> specs_pol2 = build_second_polarization_specs(specs_pol1);

    // 在 rank0 聚合两份全局解
    std::vector<std::complex<double>> sol_pol1_host;
    std::vector<std::complex<double>> sol_pol2_host;

    {
        cem::TimerScope scope_solve("solve_dual_polarizations");
        solver_manager.solve_dual_polarizations_reuse_matrix(
            em_system, specs_pol1, specs_pol2, sol_pol1_host, sol_pol2_host);
    }

    cem::PostProcessor post;
    // 基础日志仅输出网格/尺寸等
    {
        cem::TimerScope scope_post_basic("post_write_basic_log");
        post.write_basic_log(em_system);
    }

    // 频率循环：每个 color 组顺序处理其负责的若干频率

    for (int fi : freq_indices)
    {
        double freq_for_group = freq_list[fi];
        double omega_for_group = 2.0 * cem::constants::pi * freq_for_group;
        em_system.set_omega(omega_for_group);

        if (group_rank == 0)
        {
            std::cout << "[color=" << color << "][world_rank=" << world_rank
                      << "][group_rank=" << group_rank << "] "
                      << "[omega] freq_index=" << fi << ", freq=" << freq_for_group
                      << " Hz, omega=" << omega_for_group << std::endl;
            std::cout << "[color=" << color << "][world_rank=" << world_rank
                      << "][group_rank=" << group_rank << "] "
                      << "[dual] running dual polarizations with matrix reuse..." << std::endl;
        }

        // 在 rank0 聚合两份全局解（每个频率独立一份）
        std::vector<std::complex<double>> sol_pol1_host;
        std::vector<std::complex<double>> sol_pol2_host;

        {
            cem::TimerScope scope_solve("solve_dual_polarizations");
            solver_manager.solve_dual_polarizations_reuse_matrix(
                em_system, specs_pol1, specs_pol2, sol_pol1_host, sol_pol2_host);
        }

        // 输出文件名：使用频率值命名
        std::string csv_name = std::string("mt_") + std::to_string(freq_for_group) + ".csv";
        {
            cem::TimerScope scope_write_csv("post_write_mt_csv_dual");
            post.write_mt_csv_dual_from_host(
                em_system, sol_pol1_host, sol_pol2_host, params, csv_name);
        }
    }

    init.comm().barrier();
    return 0;
}